package com.test.springai.service;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

/**
 * Created with IntelliJ IDEA.
 *
 * @description: 默认的多查询扩展器实现 使用 ChatClient 生成子查询，并使用 VectorStore 并行检索文档
 * @author: liuziyang
 * @since: 9/3/25 2:12 PM
 * @modifiedBy:
 * @version: 1.0
 */
@Slf4j
@Component
public class DefaultMultiQueryExpander implements MultiQueryExpander {
  private final ChatClient chatClient;
  private final VectorStore vectorStore;

  /** 用于异步任务执行的线程池. */
  private final Executor threadPoolTaskExecutor;

  /** 用于生成子查询的LLM提示模板. 这个模板是通用的，不限制具体业务场景 */
  private final String SUB_QUERY_GENERATION_PROMPT =
      """
            你是一个能够将用户复杂问题分解为多个简单、具体的子查询的助手。
            这些子查询将用于从向量数据库中检索相关文档。
            请根据以下用户问题，生成3到5个不同的子查询，每个子查询之间用"###"分隔。
            不要包含序号或其他额外文本，只返回子查询。
            用户问题: {user_input}
            """;

  /**
   * 构造函数，用于初始化 {@code DefaultMultiQueryExpander}。
   *
   * @param chatClientBuilder 用于构建 {@code ChatClient} 实例的 {@code ChatClient.Builder}。
   * @param vectorStore 用于文档检索的 {@code VectorStore} 实例。
   * @param threadPoolTaskExecutor 用于异步任务执行的 {@code Executor} 实例（例如，一个定制的线程池）。
   */
  public DefaultMultiQueryExpander(
      ChatClient.Builder chatClientBuilder,
      VectorStore vectorStore,
      @Qualifier("applicationTaskExecutor") Executor threadPoolTaskExecutor) {
    this.chatClient = chatClientBuilder.build();
    this.vectorStore = vectorStore;
    this.threadPoolTaskExecutor = threadPoolTaskExecutor;
  }

  /**
   * 异步生成子查询。 该方法会向 LLM 发送一个提示，指导其根据原始查询生成多个相关的子查询。
   *
   * @param originalQuery 用户的原始查询字符串。
   * @return 包含子查询列表的 {@code CompletableFuture}。
   */
  @Override
  public CompletableFuture<List<String>> generateSubQueries(String originalQuery) {
    // 在指定的线程池中执行异步任务
    return CompletableFuture.supplyAsync(
        () -> {
          log.info("Generating sub-queries for: {}", originalQuery);
          SystemPromptTemplate subQueryTemplate =
              new SystemPromptTemplate(SUB_QUERY_GENERATION_PROMPT);
          // 填充模板中的 {user_input} 占位符
          String systemMessageContent =
              subQueryTemplate.createMessage(Map.of("user_input", originalQuery)).getText();
          Prompt prompt = new Prompt(List.of(new UserMessage(systemMessageContent)));
          String responseContent = chatClient.prompt(prompt).call().content();
          log.debug("Raw sub-query generation response: {}", responseContent);
          if (responseContent != null && !responseContent.trim().isEmpty()) {
            // 根据 "###" 分隔符解析子查询，并进行去空格和过滤空字符串处理
            return Stream.of(responseContent.split("###"))
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .collect(Collectors.toList());
          }
          // 如果没有生成子查询，记录警告
          log.warn("No sub-queries generated for original query: {}", originalQuery);
          return Collections.emptyList();
        },
        threadPoolTaskExecutor);
  }

  /**
   * 异步并行检索文档。 对于每个子查询，该方法会并行调用向量存储进行相似度搜索，并聚合所有结果。
   *
   * @param subQueries 子查询列表。
   * @return 包含所有检索到的文档的 {@code CompletableFuture}（可能包含重复）。
   */
  @Override
  public CompletableFuture<List<Document>> retrieveDocuments(List<String> subQueries) {
    if (subQueries == null || subQueries.isEmpty()) {
      log.info("No sub-queries provided, returning empty document list.");
      return CompletableFuture.completedFuture(Collections.emptyList());
    }
    // 使用 CompletableFuture 并行执行检索任务
    List<CompletableFuture<List<Document>>> futures =
        subQueries.stream()
            .map(
                query ->
                    CompletableFuture.supplyAsync(
                        () -> {
                          log.debug("Retrieving for sub-query: {}", query); // 记录每个子查询的检索
                          return vectorStore.similaritySearch(query); // 调用 VectorStore
                        },
                        threadPoolTaskExecutor)) // 在自定义线程池中执行异步检索
            .toList();
    // 使用 CompletableFuture.allOf 等待所有检索任务完成，然后合并结果
    return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
        .thenApply(
            v ->
                futures.stream()
                    .flatMap(
                        future -> {
                          try {
                            return future.join().stream(); // 获取异步任务的结果流
                          } catch (Exception e) {
                            log.error(
                                "Error during retrieval for a sub-query: {}",
                                e.getMessage(),
                                e); // 记录错误日志
                            return java.util.stream.Stream.empty(); // 忽略出错的检索结果
                          }
                        })
                    .collect(Collectors.toList()));
  }

  /**
   * 去重文档列表。 该方法会基于 {@code Document} 对象的 {@code equals()} 和 {@code hashCode()} 方法来识别并移除重复文档。
   *
   * @param documents 待去重的文档列表。
   * @return 去重后的文档列表。
   */
  @Override
  public List<Document> deduplicateDocuments(List<Document> documents) {
    if (documents == null || documents.isEmpty()) {
      return Collections.emptyList();
    }
    // 使用 Set 的特性自动去重 Document。
    // 这要求 Document 类正确实现 equals() 和 hashCode() 方法，
    // 通常建议基于其内容（getText()）或唯一的ID（如果有的话）来实现。
    Set<Document> uniqueDocuments = new HashSet<>(documents);
    log.debug(
        "Deduplicated documents from {} to {} unique documents.",
        documents.size(),
        uniqueDocuments.size());
    // 将 Set 转换为 List 返回
    return new ArrayList<>(uniqueDocuments);
  }
}
